#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RocketMQ 5.1.0 Python 示例代码
用于连接RocketMQ并发送消息
"""

import json
import time
from typing import Optional
from rocketmq.client import Producer, Message, PushConsumer, ConsumeStatus
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class RocketMQExample:
    """RocketMQ 示例类"""
    
    def __init__(self, namesrv_addr: str = "localhost:9876", 
                 access_key: str = "735cf4a8f557", 
                 secret_key: str = "Zjlab123!@#",
                 group_name: str = "python_example_group"):
        """
        初始化RocketMQ连接
        
        Args:
            namesrv_addr: NameServer地址
            access_key: ACL访问密钥
            secret_key: ACL密钥
            group_name: 生产者/消费者组名
        """
        self.namesrv_addr = namesrv_addr
        self.access_key = access_key
        self.secret_key = secret_key
        self.group_name = group_name
        self.producer = None
        self.consumer = None
    
    def create_producer(self) -> Producer:
        """创建生产者"""
        try:
            # 创建生产者实例
            self.producer = Producer(self.group_name)
            
            # 设置NameServer地址
            self.producer.set_name_server_address(self.namesrv_addr)
            
            # 设置ACL认证信息
            self.producer.set_session_credentials(self.access_key, self.secret_key, "ALIYUN")
            
            # 启动生产者
            self.producer.start()
            logger.info(f"生产者启动成功，NameServer: {self.namesrv_addr}")
            return self.producer
            
        except Exception as e:
            logger.error(f"创建生产者失败: {e}")
            raise
    
    def send_sync_message(self, topic: str, message_body: str, 
                         tags: Optional[str] = None, keys: Optional[str] = None) -> bool:
        """
        发送同步消息
        
        Args:
            topic: 主题名称
            message_body: 消息内容
            tags: 消息标签
            keys: 消息键
            
        Returns:
            bool: 发送是否成功
        """
        if not self.producer:
            logger.error("生产者未初始化，请先调用 create_producer()")
            return False
        
        try:
            # 创建消息对象
            msg = Message(topic)
            msg.set_body(message_body.encode('utf-8'))
            
            if tags:
                msg.set_tags(tags)
            if keys:
                msg.set_keys(keys)
            
            # 发送同步消息
            send_result = self.producer.send_sync(msg)
            logger.info(f"消息发送成功 - Topic: {topic}, MsgId: {send_result.msg_id}")
            return True
            
        except Exception as e:
            logger.error(f"发送消息失败: {e}")
            return False
    
    def send_async_message(self, topic: str, message_body: str,
                          tags: Optional[str] = None, keys: Optional[str] = None) -> bool:
        """
        发送异步消息
        
        Args:
            topic: 主题名称
            message_body: 消息内容
            tags: 消息标签
            keys: 消息键
            
        Returns:
            bool: 发送是否成功
        """
        if not self.producer:
            logger.error("生产者未初始化，请先调用 create_producer()")
            return False
        
        try:
            # 创建消息对象
            msg = Message(topic)
            msg.set_body(message_body.encode('utf-8'))
            
            if tags:
                msg.set_tags(tags)
            if keys:
                msg.set_keys(keys)
            
            # 发送异步消息
            self.producer.send_async(msg, self._send_callback)
            logger.info(f"异步消息已发送 - Topic: {topic}")
            return True
            
        except Exception as e:
            logger.error(f"发送异步消息失败: {e}")
            return False
    
    def _send_callback(self, send_result):
        """异步发送回调函数"""
        if send_result.status == 0:  # 发送成功
            logger.info(f"异步消息发送成功 - MsgId: {send_result.msg_id}")
        else:
            logger.error(f"异步消息发送失败 - Status: {send_result.status}")
    
    def create_consumer(self, topic: str, tags: str = "*") -> PushConsumer:
        """
        创建消费者
        
        Args:
            topic: 主题名称
            tags: 消息标签过滤，*表示所有标签
            
        Returns:
            PushConsumer: 消费者实例
        """
        try:
            # 创建消费者实例
            self.consumer = PushConsumer(f"{self.group_name}_consumer")
            
            # 设置NameServer地址
            self.consumer.set_name_server_address(self.namesrv_addr)
            
            # 设置ACL认证信息
            self.consumer.set_session_credentials(self.access_key, self.secret_key, "ALIYUN")
            
            # 订阅主题
            self.consumer.subscribe(topic, tags)
            
            # 注册消息监听器
            self.consumer.register_message_listener(self._message_listener)
            
            # 启动消费者
            self.consumer.start()
            logger.info(f"消费者启动成功 - Topic: {topic}, Tags: {tags}")
            return self.consumer
            
        except Exception as e:
            logger.error(f"创建消费者失败: {e}")
            raise
    
    def _message_listener(self, msg):
        """
        消息监听器
        
        Args:
            msg: 消息对象
            
        Returns:
            ConsumeStatus: 消费状态
        """
        try:
            # 解析消息内容
            message_body = msg.body.decode('utf-8')
            topic = msg.topic
            tags = msg.get_tags()
            keys = msg.get_keys()
            msg_id = msg.id
            
            logger.info(f"收到消息 - Topic: {topic}, Tags: {tags}, Keys: {keys}, MsgId: {msg_id}")
            logger.info(f"消息内容: {message_body}")
            
            # 这里可以添加您的业务逻辑
            # 例如：解析JSON、处理数据等
            
            # 返回消费成功状态
            return ConsumeStatus.CONSUME_SUCCESS
            
        except Exception as e:
            logger.error(f"处理消息失败: {e}")
            # 返回消费失败状态，消息会被重新投递
            return ConsumeStatus.RECONSUME_LATER
    
    def shutdown(self):
        """关闭连接"""
        try:
            if self.producer:
                self.producer.shutdown()
                logger.info("生产者已关闭")
            
            if self.consumer:
                self.consumer.shutdown()
                logger.info("消费者已关闭")
                
        except Exception as e:
            logger.error(f"关闭连接失败: {e}")


def main():
    """主函数示例"""
    
    # 创建RocketMQ示例实例
    # 根据您的配置修改连接参数
    mq_example = RocketMQExample(
        namesrv_addr="localhost:9876",  # 根据实际情况修改
        access_key="735cf4a8f557",
        secret_key="Zjlab123!@#",
        group_name="python_example_group"
    )
    
    try:
        # 创建生产者
        producer = mq_example.create_producer()
        
        # 定义主题名称
        topic = "test_topic"
        
        # 发送同步消息示例
        logger.info("=== 发送同步消息示例 ===")
        for i in range(3):
            message_data = {
                "id": i + 1,
                "message": f"这是第{i + 1}条同步消息",
                "timestamp": int(time.time()),
                "type": "sync"
            }
            
            success = mq_example.send_sync_message(
                topic=topic,
                message_body=json.dumps(message_data, ensure_ascii=False),
                tags="sync",
                keys=f"key_{i + 1}"
            )
            
            if success:
                logger.info(f"同步消息 {i + 1} 发送成功")
            else:
                logger.error(f"同步消息 {i + 1} 发送失败")
            
            time.sleep(1)
        
        # 发送异步消息示例
        logger.info("=== 发送异步消息示例 ===")
        for i in range(3):
            message_data = {
                "id": i + 1,
                "message": f"这是第{i + 1}条异步消息",
                "timestamp": int(time.time()),
                "type": "async"
            }
            
            success = mq_example.send_async_message(
                topic=topic,
                message_body=json.dumps(message_data, ensure_ascii=False),
                tags="async",
                keys=f"key_{i + 1}"
            )
            
            if success:
                logger.info(f"异步消息 {i + 1} 已发送")
            else:
                logger.error(f"异步消息 {i + 1} 发送失败")
            
            time.sleep(1)
        
        # 等待异步消息发送完成
        time.sleep(3)
        
        # 创建消费者示例（可选）
        logger.info("=== 创建消费者示例 ===")
        consumer = mq_example.create_consumer(topic, "*")
        
        # 保持消费者运行一段时间
        logger.info("消费者正在运行，按 Ctrl+C 停止...")
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            logger.info("收到停止信号")
        
    except Exception as e:
        logger.error(f"运行示例失败: {e}")
    
    finally:
        # 关闭连接
        mq_example.shutdown()


if __name__ == "__main__":
    main() 